# -*- coding: utf-8 -*-

'''
读取数据，返回图像和标签
label_file: 标签文件的路径，需要标签和图片文件在同一个目录下，标签的每一行代表一张图，
            一行的第一个数据代表图片的文件名，其余的是标签
count：遍历的最多的图像数，比如文件目录中，有1000张图，count=100表示只取其中的100张；<=0则表示取所有图像
patch_size: 一次放回的图像数，比如count=100， patch_size=20则表示一次从这100张图中取20张，
            合适的比较大的patch_size似乎可以提高速度；
'''

import os
import tensorflow as tf
import numpy as np
import cv2
import random

class load:
    def __init__(self, label_file, count=-1, begin_index=0):
        self.labels = open(label_file)
        self._root = os.path.dirname(label_file)
        self._patches_index = 0
        self._count = count
        self._begin_index = begin_index
        self.image_type = '.png'
        self.labels_list = []
        self.index = []
        self.current_index = 0
        self._labels_count = 0
        self.__loadLabelsInit()

    def __del__(self):
        self.labels.close()

    def __loadLabelsInit(self):
        count = self._count

        for i in range(self._begin_index):
            self.labels.readline().split()
        
        label = self.labels.readline().split()
        if count <= 0:
            while label:
                label[1:] = [float(num) for num in label[1:]]
                self.labels_list.append(label)        
                label = self.labels.readline().split()
        else:
            while (count > 0) and label:
                label[1:] = [float(num) for num in label[1:]]
                self.labels_list.append(label)        
                label = self.labels.readline().split()
                count -= 1
        self._labels_count = len(self.labels_list)
        for i in range(self._labels_count):
            self.index.append(i)


    def __load(self):
        if self._patches_index >= self._labels_count:
            self._patches_index = 0
            random.shuffle(self.index)
        
        label = self.labels_list[self.index[self._patches_index]]

        self._patches_index += 1

        image_path = self._root + '/' + str(label[0]) + self.image_type
        image = cv2.imread(image_path)

        label = [num for num in label[1:]]
    
        return image, label

    def load(self, patch_size=1):
        image, label=[], []
        if patch_size == 1:
            return self.__load()
        elif patch_size > 1:
            for i in range(patch_size):
                _image, _label = self.__load()
                image.append(_image)
                label.append(_label)
            return image, label
        else:
            exit(0)